use derive_builder::Builder;
use generator_engine::tera::TeraEngine;
use model_graph_model::model::default_model_name;
use model_graph_types::{
    assistant::prompts::{PromptTemplateItem, PromptTemplateItemBuilder, USER_ROLE_KEY},
    container::workflow::{
        LlmStatementBuilder, WorkflowBlockBody, WorkflowBlockBuilder, WorkflowBlockChain,
        WorkflowBlockChainBuilder, WorkflowBlockRef, WorkflowBlockReturnBuilder,
        WorkflowBlockReturnVariable, WorkflowBlockReturnVariableBuilder, WorkflowBlockStartBuilder,
        WorkflowBlockStatement, WorkflowLlmCompletionParamsBuilder, WorkflowLlmContextBuilder,
        WorkflowModel, WorkflowModelBuilder,
    },
    generator::CodeGenerator,
    modeling::{assistant::ModelType, Value},
};
use serde::{Deserialize, Serialize};
use std::borrow::BorrowMut;
use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};

#[derive(Debug, Serialize, Deserialize, Default, Clone, Builder, PartialEq)]
#[builder(name = "SomeChainsBuilder")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "anyhow::Error"))]
pub struct SomeChains {}

impl SomeChains {
    pub async fn get_rule_block_chain(
        &self,
        prompt_template: &String,
        instruction: &String,
    ) -> anyhow::Result<WorkflowBlockChain> {
        let engine = TeraEngine::default();
        let prompt_str = engine
            .simple(
                Value::try_from(serde_json::json!({
                    "TASK_DESCRIPTION": instruction
                }))?,
                prompt_template,
            )
            .await?;

        let (provider_name, model_name) = default_model_name(ModelType::LLM).await?;
        //构建LLM block
        let body = WorkflowBlockBody::Statement(WorkflowBlockStatement::LLM(
            LlmStatementBuilder::default()
                .model(
                    WorkflowModelBuilder::default()
                        .provider(provider_name)
                        .mode(model_name)
                        .completion_params(
                            WorkflowLlmCompletionParamsBuilder::default()
                                .temperature(0.01)
                                .build()?,
                        )
                        .build()?,
                )
                .prompt_template(vec![PromptTemplateItemBuilder::default()
                    .role(USER_ROLE_KEY)
                    .text(prompt_str)
                    .build()?])
                .context(WorkflowLlmContextBuilder::default().build()?)
                .build()?,
        ));

        Ok(self.get_single_block_chain(body, vec![String::from("text")])?)
    }

    pub async fn get_chat_block_chain(
        &self,
        prompts: Vec<PromptTemplateItem>,
        model: &WorkflowModel,
    ) -> anyhow::Result<WorkflowBlockChain> {
        //构建LLM block
        let body = WorkflowBlockBody::Statement(WorkflowBlockStatement::LLM(
            LlmStatementBuilder::default()
                .model(model.clone())
                .prompt_template(prompts)
                .context(WorkflowLlmContextBuilder::default().build()?)
                .build()?,
        ));

        Ok(self.get_single_block_chain(body, vec![String::from("text")])?)
    }

    /// 构建单节点的chain
    pub fn get_single_block_chain(
        &self,
        block: WorkflowBlockBody,
        outputs: Vec<String>,
    ) -> anyhow::Result<WorkflowBlockChain> {
        //开始
        let start_block = Arc::new(Mutex::new(
            WorkflowBlockBuilder::default()
                .id("start")
                .body(WorkflowBlockBody::Start(
                    WorkflowBlockStartBuilder::default().build()?,
                ))
                .build()?,
        ));
        //body
        let body_block = Arc::new(Mutex::new(
            WorkflowBlockBuilder::default()
                .id("body")
                .body(block)
                .build()?,
        ));
        //结束
        let end_block = Arc::new(Mutex::new(
            WorkflowBlockBuilder::default()
                .id("end")
                .body(WorkflowBlockBody::Return(
                    WorkflowBlockReturnBuilder::default()
                        .exports(
                            outputs
                                .iter()
                                .map(|v| {
                                    let variable = WorkflowBlockReturnVariableBuilder::default()
                                        .variable(v)
                                        .value_selector((String::from("body"), v.clone()))
                                        .build();
                                    match variable {
                                        Ok(item) => item,
                                        Err(_) => WorkflowBlockReturnVariable::default(),
                                    }
                                })
                                .collect::<Vec<WorkflowBlockReturnVariable>>(),
                        )
                        .build()?,
                ))
                .build()?,
        ));

        _set_outputs(&start_block, &String::from("source$target"), &body_block)?;
        _set_outputs(&body_block, &String::from("source$target"), &end_block)?;
        _set_inputs(&body_block, &String::from("source$target"), &start_block)?;
        _set_inputs(&end_block, &String::from("source$target"), &body_block)?;

        //index
        let mut index = HashMap::new();

        index.insert(String::from("start"), start_block.clone());
        index.insert(String::from("body"), body_block);
        index.insert(String::from("end"), end_block);

        Ok(WorkflowBlockChainBuilder::default()
            .root(start_block)
            .index(index)
            .build()?)
    }
}

fn _set_outputs(
    block: &WorkflowBlockRef,
    key: &String,
    target: &WorkflowBlockRef,
) -> anyhow::Result<()> {
    block
        .lock()
        .map_err(|err| anyhow::anyhow!("{}", err))?
        .borrow_mut()
        .outputs
        .insert(key.clone(), target.clone());
    Ok(())
}

fn _set_inputs(
    block: &WorkflowBlockRef,
    key: &String,
    source: &WorkflowBlockRef,
) -> anyhow::Result<()> {
    block
        .lock()
        .map_err(|err| anyhow::anyhow!("{}", err))?
        .borrow_mut()
        .inputs
        .insert(key.clone(), source.clone());
    Ok(())
}
